from __future__ import annotations

import functools
import os
from collections.abc import Callable
from copy import deepcopy
from typing import Any

import torch.distributed as dist
from datasets import Dataset
from torchdata.stateful_dataloader import StatefulDataLoader

from areal.api.alloc_mode import AllocationMode
from areal.api.cli_args import (
    InferenceEngineConfig,
    PPOActorConfig,
    PPOConfig,
    PPOCriticConfig,
    TrainDatasetConfig,
    ValidDatasetConfig,
)
from areal.api.engine_api import InferenceEngine
from areal.api.io_struct import FinetuneSpec, StepInfo, WeightUpdateMeta
from areal.api.workflow_api import RolloutWorkflow
from areal.engine.megatron_engine import MegatronEngine
from areal.engine.ppo.actor import FSDPPPOActor, MegatronPPOActor
from areal.engine.ppo.critic import FSDPPPOCritic, MegatronPPOCritic
from areal.engine.sglang_remote import RemoteSGLangEngine
from areal.engine.vllm_remote import RemotevLLMEngine
from areal.platforms import current_platform
from areal.utils import logging, perf_tracer, seeding, stats_tracker
from areal.utils.dataloader import create_dataloader
from areal.utils.device import log_gpu_stats
from areal.utils.evaluator import Evaluator
from areal.utils.hf_utils import load_hf_processor_and_tokenizer
from areal.utils.perf_tracer import Category
from areal.utils.recover import RecoverHandler
from areal.utils.saver import Saver
from areal.utils.stats_logger import StatsLogger

logger = logging.getLogger(__name__)


class PPOTrainer:
    def __init__(
        self,
        config: PPOConfig,
        train_dataset: Dataset,
        valid_dataset: Dataset | None = None,
    ):
        rank = int(os.getenv("RANK", "0"))
        # Configure performance tracer
        if config.perf_tracer is not None:
            perf_tracer.configure(config.perf_tracer, rank=rank)

        self.config = config
        self.processor, self.tokenizer = load_hf_processor_and_tokenizer(
            config.tokenizer_path
        )

        # Set seed.
        seeding.set_random_seed(config.seed, key=f"trainer{rank}")

        # Parse allocation mode.
        self.allocation_mode = AllocationMode.from_str(config.allocation_mode)

        # Create models: actor, critic, etc.
        self.actor = self._create_actor(config.actor)
        self.critic = None
        if config.critic is not None:
            self.critic = self._create_critic(config.critic)
        self.ref = None
        if config.actor.kl_ctl > 0 and config.ref is not None:
            self.ref = self._create_actor(config.ref)

        # Create dataloaders
        self.train_dataset = train_dataset
        self.valid_dataset = valid_dataset
        self.train_dataloader = self._create_dataloader(
            train_dataset,
            dataset_config=self.config.train_dataset,
            rank=self.actor.data_parallel_rank,
            world_size=self.actor.data_parallel_world_size,
        )
        self.valid_dataloader = None
        if self.config.valid_dataset is not None and valid_dataset is not None:
            self.valid_dataloader = self._create_dataloader(
                valid_dataset,
                dataset_config=self.config.valid_dataset,
                rank=self.actor.data_parallel_rank,
                world_size=self.actor.data_parallel_world_size,
            )

        # Initialize inference
        self.rollout = self._init_rollout(config.rollout, is_eval=False)
        self.eval_rollout = self._init_rollout(config.rollout, is_eval=True)

        ft_spec = FinetuneSpec(
            total_train_epochs=config.total_train_epochs,
            dataset_size=len(self.train_dataloader) * config.train_dataset.batch_size,
            train_batch_size=config.train_dataset.batch_size,
        )

        # Initialize models
        self.parallel_strategy = self.allocation_mode.train
        assert self.parallel_strategy is not None
        engine_init_kwargs = {"addr": None, "ft_spec": ft_spec}
        self.actor.initialize(**engine_init_kwargs)
        if self.critic is not None:
            self.critic.initialize(**engine_init_kwargs)
        if self.ref is not None:
            self.ref.initialize(**engine_init_kwargs)

        # Prepare weight update meta and connect to inference engine
        if self.config.actor.weight_update_mode == "disk":
            self.weight_update_meta = WeightUpdateMeta.from_disk(
                experiment_name=config.experiment_name,
                trial_name=config.trial_name,
                file_root=config.cluster.fileroot,
                name="default",
                use_lora=config.actor.use_lora,
                clear_checkpoint_after_load=True,
            )
        elif self.config.actor.weight_update_mode == "xccl":
            # NCCL/XCCL weight update
            if isinstance(self.actor, MegatronEngine):
                self.weight_update_meta = WeightUpdateMeta.from_megatron_xccl(
                    self.allocation_mode,
                    nccl_group_name=self.actor.weight_update_group_name,
                )
            else:
                self.weight_update_meta = WeightUpdateMeta.from_fsdp_xccl(
                    self.allocation_mode
                )
        else:
            raise ValueError(
                f"Invalid weight update mode: {self.config.actor.weight_update_mode}"
            )
        self.actor.connect_engine(self.rollout, self.weight_update_meta)

        # Set up evaluation
        self.evaluator = Evaluator(config.evaluator, ft_spec)

        # Set up save as HF model
        self.saver = Saver(config.saver, ft_spec)
        self.recover_handler = RecoverHandler(config.recover, ft_spec)

        # Set up statistics logging (wandb, tensoboard, etc.)
        self.stats_logger = StatsLogger(config, ft_spec)

        # Set up checkpointing for recover
        self.recover_info = self.recover_handler.load(
            self.actor,
            self.saver,
            self.evaluator,
            self.stats_logger,
            self.train_dataloader,
            inference_engine=self.rollout,
            weight_update_meta=self.weight_update_meta,
        )

    def train(
        self,
        workflow: RolloutWorkflow,
        eval_workflow: RolloutWorkflow | None = None,
        dynamic_filter_fn: Callable[[dict[str, Any]], bool] | str | None = None,
    ):
        config = self.config
        start_step = (
            self.recover_info.last_step_info.next().global_step
            if self.recover_info is not None
            else 0
        )

        total_epochs = config.total_train_epochs
        steps_per_epoch = len(self.train_dataloader)
        max_steps = total_epochs * steps_per_epoch

        for global_step in range(start_step, max_steps):
            if (
                config.total_train_steps is not None
                and global_step >= config.total_train_steps
            ):
                break
            epoch = global_step // steps_per_epoch
            step = global_step % steps_per_epoch

            with (
                stats_tracker.record_timing("rollout"),
                perf_tracer.trace_scope(
                    "train.rollout",
                    category=Category.COMPUTE,
                    args={
                        "global_step": global_step,
                        "epoch_step": step,
                    },
                ),
            ):
                batch = self.actor.prepare_batch(
                    self.train_dataloader,
                    granularity=self.config.actor.group_size,
                    workflow=workflow,
                    should_accept_fn=dynamic_filter_fn,
                )

            if self.critic is not None:
                with (
                    stats_tracker.record_timing("critic_values"),
                    perf_tracer.trace_scope(
                        "train.compute_values",
                        category=Category.COMPUTE,
                        args={"global_step": global_step},
                    ),
                ):
                    values = self.critic.compute_values(batch)
                    batch["values"] = values
                    log_gpu_stats("critic values")

            if config.actor.recompute_logprob or config.actor.use_decoupled_loss:
                with (
                    stats_tracker.record_timing("recompute_logp"),
                    perf_tracer.trace_scope(
                        "train.recompute_logp",
                        category=Category.COMPUTE,
                        args={"global_step": global_step},
                    ),
                ):
                    logp = self.actor.compute_logp(batch)
                    batch["prox_logp"] = logp
                    log_gpu_stats("recompute logp")

            if self.ref is not None:
                with (
                    stats_tracker.record_timing("ref_logp"),
                    perf_tracer.trace_scope(
                        "train.ref_logp",
                        category=Category.COMPUTE,
                        args={"global_step": global_step},
                    ),
                ):
                    batch["ref_logp"] = self.ref.compute_logp(batch)
                    log_gpu_stats("ref logp")

            with (
                stats_tracker.record_timing("compute_advantage"),
                perf_tracer.trace_scope(
                    "train.compute_advantage",
                    category=Category.COMPUTE,
                    args={"global_step": global_step},
                ),
            ):
                self.actor.compute_advantages(batch)
                log_gpu_stats("compute advantages")

            with (
                stats_tracker.record_timing("train_step"),
                perf_tracer.trace_scope(
                    "train.ppo_update",
                    category=Category.COMPUTE,
                    args={"global_step": global_step},
                ),
            ):
                self.actor.ppo_update(batch)
                self.actor.step_lr_scheduler()
                log_gpu_stats("ppo update")

            if self.critic is not None:
                with (
                    stats_tracker.record_timing("critic_train_step"),
                    perf_tracer.trace_scope(
                        "train.critic_ppo_update",
                        category=Category.COMPUTE,
                        args={"global_step": global_step},
                    ),
                ):
                    self.critic.ppo_update(batch)
                    self.critic.step_lr_scheduler()
                    log_gpu_stats("ppo critic update")

            # pause inference for updating weights, save, and evaluation
            self.rollout.pause()

            with (
                stats_tracker.record_timing("update_weights"),
                perf_tracer.trace_scope(
                    "train.update_weights",
                    category=Category.COMM,
                    args={"global_step": global_step},
                ),
            ):
                self.actor.update_weights(self.weight_update_meta)

                self.actor.set_version(global_step + 1)
                if self.critic is not None:
                    self.critic.set_version(global_step + 1)
                self.rollout.set_version(global_step + 1)
                self.eval_rollout.set_version(global_step + 1)

            with (
                stats_tracker.record_timing("save"),
                perf_tracer.trace_scope(
                    "train.save",
                    category=Category.IO,
                    args={"global_step": global_step},
                ),
            ):
                self._save_hf(epoch=epoch, epoch_step=step, global_step=global_step)

            with (
                stats_tracker.record_timing("checkpoint_for_recover"),
                perf_tracer.trace_scope(
                    "train.checkpoint",
                    category=Category.IO,
                    args={"global_step": global_step},
                ),
            ):
                self._save_recover_checkpoint(
                    epoch=epoch, epoch_step=step, global_step=global_step
                )

            with (
                stats_tracker.record_timing("eval"),
                perf_tracer.trace_scope(
                    "train.eval",
                    category=Category.COMPUTE,
                    args={"global_step": global_step},
                ),
            ):
                self._evaluate(
                    eval_workflow=eval_workflow,
                    epoch=epoch,
                    epoch_step=step,
                    global_step=global_step,
                )

            with perf_tracer.trace_scope(
                "train.log_stats",
                category=Category.INSTR,
                args={"global_step": global_step},
            ):
                self._export_and_commit_stats(
                    epoch=epoch, epoch_step=step, global_step=global_step
                )

            # Resume rollout
            self.rollout.resume()

            perf_tracer.save(step=global_step)

    def close(self):
        self.stats_logger.close()
        self.eval_rollout.destroy()
        self.rollout.destroy()
        if self.ref is not None:
            self.ref.destroy()
        if self.critic is not None:
            self.critic.destroy()
        self.actor.destroy()
        perf_tracer.save(force=True)

    def _create_dataloader(
        self,
        dataset: Dataset,
        dataset_config: TrainDatasetConfig | ValidDatasetConfig,
        rank: int,
        world_size: int,
    ) -> StatefulDataLoader:
        return create_dataloader(
            dataset,
            rank=rank,
            world_size=world_size,
            dataset_config=dataset_config,
        )

    def _create_actor(self, actor_config: PPOActorConfig):
        if self.allocation_mode.train_backend == "fsdp":
            actor = FSDPPPOActor(config=actor_config)
        elif self.allocation_mode.train_backend == "megatron":
            actor = MegatronPPOActor(config=actor_config)
        else:
            raise ValueError(
                f"Invalid backend: {self.allocation_mode.train_backend}, expected fsdp or megatron"
            )
        actor.create_process_group(parallel_strategy=self.allocation_mode.train)
        return actor

    def _create_critic(self, critic_config: PPOCriticConfig):
        if self.allocation_mode.train_backend == "fsdp":
            critic = FSDPPPOCritic(config=critic_config)
        elif self.allocation_mode.train_backend == "megatron":
            critic = MegatronPPOCritic(config=critic_config)
        else:
            raise ValueError(
                f"Invalid backend: {self.allocation_mode.train_backend}, expected fsdp or megatron"
            )
        critic.create_process_group(parallel_strategy=self.allocation_mode.train)
        return critic

    def _init_rollout(
        self, rollout_config: InferenceEngineConfig, is_eval: bool = False
    ) -> InferenceEngine:
        # Initialize inference engine
        if self.allocation_mode.gen_backend == "sglang":
            engine = RemoteSGLangEngine(deepcopy(rollout_config))
        elif self.allocation_mode.gen_backend == "vllm":
            engine = RemotevLLMEngine(deepcopy(rollout_config))
        else:
            raise ValueError(
                f"Invalid backend: {self.allocation_mode.gen_backend}, expected sglang or vllm"
            )

        if is_eval:
            # NOTE: eval does not have any offpolicyness control
            engine.config.max_head_offpolicyness = int(1e12)
        engine.initialize(train_data_parallel_size=self.allocation_mode.train.dp_size)
        return engine

    def _save_hf(self, epoch: int, epoch_step: int, global_step: int):
        # Save as HF models for evaluation
        self.saver.save(
            self.actor,
            epoch,
            epoch_step,
            global_step,
            tokenizer=self.tokenizer,
            processor=self.processor,
        )
        if self.critic is not None:
            self.saver.save(
                self.critic,
                epoch,
                epoch_step,
                global_step,
                tokenizer=self.tokenizer,
                processor=self.processor,
                name="critic",
            )
        dist.barrier(device_ids=[self.actor.device.index])
        current_platform.synchronize()

    def _save_recover_checkpoint(self, epoch: int, epoch_step: int, global_step: int):
        # Save recoverable checkpoints
        to_save = dict(default=self.actor)
        if self.critic is not None:
            to_save["critic"] = self.critic
        step_info = StepInfo(
            global_step=global_step,
            epoch=epoch,
            epoch_step=epoch_step,
            steps_per_epoch=len(self.train_dataloader),
        )
        self.recover_handler.dump(
            to_save,
            step_info,
            self.saver,
            self.evaluator,
            self.stats_logger,
            self.train_dataloader,
            tokenizer=self.tokenizer,
            processor=self.processor,
        )

        dist.barrier(device_ids=[self.actor.device.index])
        current_platform.synchronize()

    def _evaluate_fn(self, eval_workflow: RolloutWorkflow):
        if self.actor.is_data_parallel_head():
            cnt = 0
            for data in self.valid_dataloader:
                for item in data:
                    self.eval_rollout.submit(item, eval_workflow)
                    cnt += 1
            self.eval_rollout.wait(cnt, timeout=None)
        dist.barrier(device_ids=[self.actor.device.index])
        current_platform.synchronize()

    def _evaluate(
        self,
        eval_workflow: RolloutWorkflow | None,
        epoch: int,
        epoch_step: int,
        global_step: int,
    ):
        if self.valid_dataloader is None or eval_workflow is None:
            return
        self.evaluator.evaluate(
            functools.partial(self._evaluate_fn, eval_workflow=eval_workflow),
            epoch,
            epoch_step,
            global_step,
        )
        dist.barrier(device_ids=[self.actor.device.index])
        current_platform.synchronize()

    def _export_and_commit_stats(self, epoch: int, epoch_step: int, global_step: int):
        # Upload statistics to the logger (e.g., wandb)
        stats = self.actor.export_stats()
        self.stats_logger.commit(epoch, epoch_step, global_step, stats)

        dist.barrier(device_ids=[self.actor.device.index])
        current_platform.synchronize()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()
        if exc_type is not None:
            raise exc_value
